import os
import time

import faiss
import numpy

numpy.random.seed(0) #设置随机数种子，保证每次运行结果一致
vectors = numpy.random.rand(1000000, 768)  # 创建1000000个随机向量 , 768表示向量的维度
query = numpy.random.rand(1, 768) # 创建1个随机向量 ,用来模拟查询向量

def fun_1():
    pass
    # 1 在CPU上创建索引
    index = faiss.IndexFlatL2(768)
    # 2 将索引移动到GPU上

    #第一个参数： GPU的计算对象
    #第二个参数： GPU的编号
    #第三个参数： 索引
    index_gpu = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, index)
    index_gpu.add(vectors)
    # 3 在GPU上检索
    start = time.time()
    D, I = index_gpu.search(query, k=2)
    print(f"普通线性索引耗时：{time.time() - start}")
    print(f"最相似的向量索引：{I}")
    print(f"最相似的向量距离：{D}")